import streamlit as st
import random
import json
from streamlit_tags import st_tags
import jsonlines

# Basic Streamlit page configuration
st.set_page_config(page_title='民事案件类案检索系统', layout='wide')

# Load predicate descriptions from predicate.json
with open('predicate.json', 'r', encoding='utf-8') as f:
    predicate_list = json.load(f)

# Load candidate case data from candidate.json
with open('candidate.json', 'r', encoding='utf-8') as f:
    candidate_cases = json.load(f)

# Load entity data from entity.jsonl
entity_cases = []
with jsonlines.open('entity.jsonl', 'r') as reader:
    for obj in reader:
        entity_cases.append(obj)

# Initialize session state for analysis trigger, selected case, analysis counter, and case description
if 'analyze_triggered' not in st.session_state:
    st.session_state['analyze_triggered'] = False
if 'selected_case' not in st.session_state:
    st.session_state['selected_case'] = None
if 'analysis_counter' not in st.session_state:
    st.session_state['analysis_counter'] = 0
if 'case_description' not in st.session_state:
    st.session_state['case_description'] = ''

# Page Title
# st.title('民事案件类案检索系统')

# Main Layout: Sidebar and Main Content with Spacing Adjustment
sidebar, main_content = st.columns([2, 3], gap='large')

# Sidebar Navigation
with sidebar:
    st.header('案件信息')
    # Text area bound to session state
    case_description = st.text_area(
        '案件描述',
        height=200,
        placeholder='请输入案件描述...',
        value=st.session_state['case_description'],
        key='case_description_input'
    )
    if st.button('分析案件'):
        # Update case description in session state
        st.session_state['case_description'] = case_description
        if case_description.strip():
            # Increment analysis counter to ensure unique st_tags key
            st.session_state['analysis_counter'] += 1
            # Randomly select a case from entity.jsonl
            st.session_state['selected_case'] = random.choice(entity_cases)
            # Update case description with selected case content from candidate.json
            case_index = st.session_state['selected_case']['index']
            if case_index in candidate_cases:
                st.session_state['case_description'] = "\n".join(candidate_cases[case_index]['content'])
            else:
                st.session_state['case_description'] = "未找到对应案例内容"
            st.session_state['analyze_triggered'] = True

# Main Content (Only show after analysis)
if st.session_state['analyze_triggered'] and st.session_state['selected_case']:
    with sidebar:
        # Extract event types and entities from the selected case
        selected_case = st.session_state['selected_case']
        event_types = list(selected_case['atom'].keys())  # List all event types, e.g., ["借款事件", ...]

        # Display detected event types with dynamic key
        st.subheader('检测到的法律事件类型')
        st_tags(
            label='',
            value=event_types,
            suggestions=[],  # Disable user input
            maxtags=5,
            key=f'event_type_tags_{st.session_state["analysis_counter"]}',  # Dynamic key
            text=""
        )
        
        # Display recognized entities for all event types, grouped by event type in expanders
        st.subheader('识别出的法律实体')
        for event_type in event_types:
            with st.expander(f"{event_type}"):
                entities = selected_case['atom'][event_type]  # Dictionary of entity roles and values for this event
                for role, value in entities.items():
                    if value != "不存在":  # Filter out entities with value "不存在"
                        st.markdown(f"""
                        <div style='padding: 12px; background-color: #f8f9fa; border-left: 6px solid #007bff; border-radius: 8px; margin-bottom: 12px;'>
                            <strong style='color: #007bff;'>{role}</strong><br>
                            <span style='color: #333;'>{value}</span>
                        </div>
                        """, unsafe_allow_html=True)

        # Extracted Legal Elements dynamically from candidate.json
        st.subheader('抽取出的法律要件')
        case_index = selected_case['index']  # Get the index from entity.jsonl
        if case_index in candidate_cases:
            case_data = candidate_cases[case_index]
            legal_elements = [predicate_list[i] for i, v in enumerate(case_data['feature']) if v == 1]
        else:
            legal_elements = []  # Fallback if index not found
            st.warning(f"未找到索引 {case_index} 对应的法律要件数据")
        st.markdown('<ul style="list-style-type: none; padding-left: 0;">', unsafe_allow_html=True)
        for element in legal_elements:
            st.markdown(f"<li style='padding: 10px; background-color: #ffe4e1; border-radius: 8px; margin-bottom: 8px;'><strong>{element}</strong></li>", unsafe_allow_html=True)
        st.markdown('</ul>', unsafe_allow_html=True)

    with main_content:
        st.header('类案检索结果')
        st.write('以下是与当前案件最相似的5个案例：')
        # Select random candidate cases
        candidate_keys = random.sample(list(candidate_cases.keys()), 10)
        for key in candidate_keys:
            case = candidate_cases[key]
            title = f"{case['caseID']}"
            score = round(random.uniform(0.7, 0.99), 2)  # Simulated similarity score
            legal_components = [predicate_list[i] for i, v in enumerate(case['feature']) if v == 1]
            related_components = [comp for comp in legal_components if comp in legal_elements]
            with st.expander(f"{title} (相似度: {score})"):
                st.markdown("**案例详情：**")
                st.markdown("""<div style='white-space: pre-wrap;'>
                {}
                </div>""".format("\n".join(case['content'])), unsafe_allow_html=True)
                st.markdown('<ul style="list-style-type: none; padding-left: 0;">', unsafe_allow_html=True)
                for comp in legal_components:
                    if comp in related_components:
                        st.markdown(f"<li style='padding: 10px; background-color: #ffe4e1; border-radius: 8px; margin-bottom: 8px;'><strong>{comp}</strong> (相关)</li>", unsafe_allow_html=True)
                    else:
                        st.markdown(f"<li style='padding: 10px; background-color: #e8f4f8; border-radius: 8px; margin-bottom: 8px;'>{comp}</li>", unsafe_allow_html=True)
                st.markdown('</ul>', unsafe_allow_html=True)
else:
    with main_content:
        st.info('请输入案件描述并点击 **分析案件** 以查看分析结果。')